import numpy as np  
import matplotlib.pyplot as plt  
from filterpy.kalman import ExtendedKalmanFilter as EKF  
from scipy.spatial.transform import Rotation as R  
  
# 生成螺旋上升运动的三维坐标数据  
def generate_spiral():  
    """Generate coordinates for a spiral motion."""  
    coordinates = []  
    t = np.linspace(0, 10, int(10 * 30), endpoint=False)  # 0到10秒，30Hz  
    for ti in t:  
        # 螺旋线的参数方程  
        x = amplitude * np.cos(ti)  
        y = amplitude * np.sin(ti)  
        z = amplitude * ti  # 时间作为z轴（上升） 
        coordinates.append([x, y, z])  
    return np.array(coordinates)  
  
# 添加高斯噪声  
def add_gaussian_noise(data, noise_sigma=10):  
    """Add Gaussian noise to the data."""  
    noisy_data = data + np.random.normal(0, noise_sigma, data.shape)  
    return noisy_data  
  
# 扩展卡尔曼滤波器模型定义  
def f(x, dt):  
    """State transition function for nearly constant velocity."""  
    return np.array([x[0] + x[3]*dt, x[1] + x[4]*dt, x[2] + x[5]*dt, x[3], x[4], x[5]])   
  
def h(x):  
    """Measurement function."""
    return np.array([x[0], x[1], x[2]])  
  
def F(x, dt):  
    """State transition jacobian."""  

    return  np.array([1, 0, 0, dt, 0, 0],  
        [0, 1, 0, 0, dt, 0],  
        [0, 0, 1, 0, 0, dt],  
        [0, 0, 0, 1, 0, 0],  
        [0, 0, 0, 0, 1, 0],  
        [0, 0, 0, 0, 0, 1] )
  
def H(x):  
    """Measurement jacobian."""  
    return np.array([[1, 0, 0, 0, 0, 0],  
                     [0, 1, 0, 0, 0, 0],  
                     [0, 0, 1, 0, 0, 0]])  


# 设置时间、频率和振幅参数  
t_total = 10  # 总时间（秒）  
frequency = 30  # 频率（Hz）  
amplitude = 100  # 振幅（厘米），为了看到明显的螺旋线，这里选择较大的振幅值  
dt = 1 / frequency  # 时间步长（秒）  
num_points = int(t_total / dt)  # 总点数  
times = np.arange(0, t_total, dt)  # 时间数组  
  
# 生成真实数据（无噪声）  
#true_data = generate_spiral(t_total, frequency, amplitude=amplitude, dt=dt) * 100 
true_data = generate_spiral() 
# 添加噪声以模拟摄像头数据  
noisy_data = add_gaussian_noise(true_data, noise_sigma=10)  # 噪声标准差为10cm（转换为厘米）
print(true_data[1])  
lidar_data = add_gaussian_noise(true_data, noise_sigma=0.1)      
print(lidar_data[1])
print(np.linalg.norm(true_data[1]-lidar_data[1]))
qconstant = 1000
qconstant1 =1000
rconstant = 1
rconstant1=0
# sequence
# 设置扩展卡尔曼滤波器参数并初始化滤波器状态向量和协方差矩阵  
ekf = EKF(dim_x=6, dim_z=3)  # 状态向量为[x, y, z]，测量向量为[x, y, z]  
ekf.x = np.array([noisy_data[0,0], noisy_data[0,1], noisy_data[0,2], 0, 0, 0])  # 初始状态（位置和速度）设置为第一个测量值，速度为0  
ekf.P = np.eye(6) * 1**2  # 初始状态协方差矩阵，对角线上的值表示对应状态的不确定性程度（这里假设所有状态的不确定性相同）  

ekf.R = np.eye(3) * (10*rconstant)**2   # 测量噪声协方差矩阵，设置为噪声的标准差的平方（转换为厘米） 
ekf.Q = np.eye(6) * (qconstant**2)  # 过程噪声协方差矩阵，这里假设过程噪声较小，因此设置为较小的值（可以根据实际情况调整）  
ekf.F = np.array([[1, 0, 0, dt, 0, 0],  
                  [0, 1, 0, 0, dt, 0],  
                  [0, 0, 1, 0, 0, dt],  
                  [0, 0, 0, 1, 0, 0],  
                  [0, 0, 0, 0, 1, 0],  
                  [0, 0, 0, 0, 0, 1]]) # 状态转移函数的雅可比矩阵（使用lambda函数传递dt）  
ekf.H = H  # 测量函数的雅可比矩阵（不需要状态向量，因此直接传递函数）  
ekf.f = f  # 状态转移函数（使用lambda函数传递dt）  
ekf.h = h  # 测量函数（不需要状态向量，因此直接传递函数）  
#ekf.R[:] = (10*100)**2  # 更新测量噪声协方差矩阵的所有元素为相同值（因为我们认为所有方向的测量噪声都相同）
ekf1 = EKF(dim_x=6, dim_z=3)  # 状态向量为[x, y, z]，测量向量为[x, y, z]  
ekf1.x = np.array([noisy_data[0,0], noisy_data[0,1], noisy_data[0,2], 0, 0, 0])  # 初始状态（位置和速度）设置为第一个测量值，速度为0  
ekf1.P = np.eye(6) * 1**2  # 初始状态协方差矩阵，对角线上的值表示对应状态的不确定性程度（这里假设所有状态的不确定性相同）  

ekf1.R = np.eye(3) * (rconstant1)**2   # 测量噪声协方差矩阵，设置为噪声的标准差的平方（转换为厘米） 
ekf1.Q = np.eye(6) * (qconstant1**2)  # 过程噪声协方差矩阵，这里假设过程噪声较小，因此设置为较小的值（可以根据实际情况调整）  
ekf1.F = np.array([[1, 0, 0, dt, 0, 0],  
                  [0, 1, 0, 0, dt, 0],  
                  [0, 0, 1, 0, 0, dt],  
                  [0, 0, 0, 1, 0, 0],  
                  [0, 0, 0, 0, 1, 0],  
                  [0, 0, 0, 0, 0, 1]]) # 状态转移函数的雅可比矩阵（使用lambda函数传递dt）  
ekf1.H = H  # 测量函数的雅可比矩阵（不需要状态向量，因此直接传递函数）  
ekf1.f = f  # 状态转移函数（使用lambda函数传递dt）  
ekf1.h = h  # 测量函数（不需要状态向量，因此直接传递函数）  

filtered_data_nom = []  # 用于存储滤波后的数据  
filtered_data_seq = []  # 用于存储滤波后的数据 
delt_nom = []
delt_seq = []
delt_lid=[]
token = 0

for tt in noisy_data:  # 对每个测量值进行滤波处理  
    ekf.update(tt,HJacobian=ekf.H,Hx=ekf.h)
      # 使用当前测量值更新滤波器状态估计和协方差矩阵  
    ekf.predict()  # 使用状态转移函数预测下一个时间步的状态估计和协方差矩阵（尽管这里我们实际上是在处理当前时间步的下一个测量值）
    filtered_data_nom.append(ekf.x)  # 存储滤波后的位置估计（忽略速度估计）  
    delt_nom.append(np.linalg.norm(ekf.x[:3] - true_data[token]))
    #print(token)
    token = token+1
filtered_data_nom = np.array(filtered_data_nom)  # 将列表转换为NumPy数组以便于绘图处理  
delt_nom = np.array(delt_nom)

token = 0
for tt in lidar_data:  # 对每个测量值进行滤波处理  
    #if(token%3==0):
    #    ekf1.update(lidar_data[token],HJacobian=ekf1.H,Hx=ekf1.h)
        # 使用当前测量值更新滤波器状态估计和协方差矩阵  
    #    ekf1.predict() 
    
    ekf1.update(tt,HJacobian=ekf1.H,Hx=ekf1.h)
    # 使用当前测量值更新滤波器状态估计和协方差矩阵  
    ekf1.predict()  # 使用状态转移函数预测下一个时间步的状态估计和协方差矩阵（尽管这里我们实际上是在处理当前时间步的下一个测量值）
    delt_seq.append(np.linalg.norm(ekf1.x[:3] - true_data[token]))
    delt_lid.append(np.linalg.norm(lidar_data[token] - true_data[token]))
    print("np.linalg.norm(lidar_data[token] - true_data[token])")
    print(np.linalg.norm(lidar_data[token] - true_data[token]))
    token = token+1

    filtered_data_seq.append(ekf1.x)  # 存储滤波后的位置估计（忽略速度估计）  
filtered_data_seq = np.array(filtered_data_seq)  # 将列表转换为NumPy数组以便于绘图处理  
delt_seq = np.array(delt_seq)
delt_lid = np.array(delt_lid)


fig = plt.figure()  


# Cam
ax1 = fig.add_subplot(221, projection='3d')  
ax1.plot(filtered_data_nom[:, 0], filtered_data_nom[:, 1], filtered_data_nom[:, 2], label='Original')  # 原图为蓝色  
ax1.set_xlabel("X Position (cm)")  
ax1.set_ylabel("Y Position (cm)")  
ax1.set_zlabel("Z Position (cm)")  
ax1.set_title("Trajectory with noise")  
plt.legend()  


# sequence
ax2 = fig.add_subplot(222, projection='3d')  
ax2.plot(filtered_data_seq[:, 0] , filtered_data_seq[:, 1] , filtered_data_seq[:, 2] , 'g', label='With Noise')  # 带噪音的图为粉红色  
ax2.set_xlabel("X Position (cm)")  
ax2.set_ylabel("Y Position (cm)")  
ax2.set_zlabel("Z Position (cm)")  
ax2.set_title("Trajectory after sequence EKF Filter")  
plt.legend() 

t_ = np.linspace(0, 10, int(10 * 30), endpoint=False)
# Cam 

ax3 = fig.add_subplot(223)  
ax3.plot(t_, np.sqrt(delt_nom), label='nom_residual')  
ax3.plot(t_ , np.sqrt(delt_seq) ,  'pink', label='seq_residual')  
print(np.sqrt(delt_seq)[1])
print(delt_seq[1])
print(filtered_data_seq[1][:3]-noisy_data[1])
ax3.plot(t_, np.mean(np.sqrt(delt_nom)) * np.ones(int(10 * 30)), linestyle='--', color='pink',label='nom_mean')
ax3.plot(t_ , np.mean(np.sqrt(delt_seq))*np.ones(int(10 * 30)) ,linestyle='--', color='grey', label='seq_mean')    

ax3.set_xlabel("Step (33.3ms)")  
ax3.set_ylabel("Residual (cm)")  
ax3.set_title("Residuals")  
ax3.set_ylim(0,10)   
plt.legend()  


# sequence Conv
ax4 = fig.add_subplot(224)  
ax4.plot(t_ , np.sqrt(delt_seq) ,  'pink', label='seq residual')  
ax4.plot(t_ , np.mean(np.sqrt(delt_seq))*np.ones(int(10 * 30)) ,'grey',   label='mean')    
ax4.plot(t_ , np.sqrt(delt_lid) ,  'grey', label='seq residual')  

ax4.set_xlabel("Step (33.3ms)")  
ax4.set_ylabel("Residual (cm)")  
ax4.set_title("Residuals after sequence EKF Filter")  
ax4.set_ylim(0,2.3)
plt.legend() 



plt.show()
